import multiprocessing as mp
from collections import deque

import gymnasium as gym
import numpy as np
import pygame
from gymnasium import spaces

from .watermelon import WatermelonGame


class WatermelonEnv(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 960}

    def __init__(
        self,
        render_mode=None,
        width=400,
        height=600,
        discrete_positions=20,
        frame_skip=1,
        stack_frames=4,  # 新增参数：堆叠帧数
    ):
        super().__init__()

        # 离散化屏幕宽度为discrete_positions个位置，再加上"不投放"动作(0号)
        self.discrete_positions = discrete_positions

        # 帧跳过设置 - 每个动作执行多少帧的物理更新
        self.frame_skip = frame_skip

        # 帧堆叠设置
        self.stack_frames = stack_frames

        # 动作空间: 0表示不投放，1-discrete_positions表示在不同位置投放
        self.action_space = spaces.Discrete(discrete_positions + 1)

        # 修改观察空间，通道数乘以堆叠帧数
        self.observation_space = spaces.Box(
            low=0, high=255, shape=(height, width, 3 * stack_frames), dtype=np.uint8
        )

        # 渲染设置
        self.render_mode = render_mode
        self.width = width
        self.height = height

        # 创建游戏实例
        self.game = WatermelonGame(width, height)

        # 初始化帧缓冲
        self.frames_buffer = deque(maxlen=stack_frames)

        # 用于人类渲染的窗口
        self.window = None
        self.clock = None

        # 记录上一次得分和动作
        self._last_score = 0
        self._last_action = 0

    def _get_observation(self):
        """获取当前游戏画面作为观察，并处理帧堆叠"""
        screen = self.game.render()
        current_frame = pygame.surfarray.array3d(screen).transpose(1, 0, 2)

        # 如果帧缓冲区不满，用当前帧填充
        while len(self.frames_buffer) < self.stack_frames:
            self.frames_buffer.append(current_frame)

        # 添加当前帧到缓冲区
        self.frames_buffer.append(current_frame)

        # 堆叠帧（在通道维度上连接）
        stacked_frames = np.concatenate(list(self.frames_buffer), axis=2)

        return stacked_frames

    def reset(self, seed=None, options=None):
        """重置环境"""
        super().reset(seed=seed)

        # 重置游戏
        self.game.reset()

        # 清空帧缓冲区
        self.frames_buffer.clear()

        # 获取初始帧
        initial_frame = pygame.surfarray.array3d(self.game.render()).transpose(1, 0, 2)

        # 用初始帧填充缓冲区
        for _ in range(self.stack_frames):
            self.frames_buffer.append(initial_frame)

        # 获取堆叠后的观察
        observation = self._get_observation()
        info = {}

        # 渲染初始帧
        if self.render_mode == "human":
            self._render_frame()

        # 重置得分和动作记录
        self._last_score = 0
        self._last_action = 0

        return observation, info

    def step(self, action):
        """执行一步动作并跳过中间帧"""
        game_over = False
        victory = False
        cumulative_reward = 0.0

        # 处理第一帧的动作
        if action is not None:
            if action == 0:
                # 特殊动作0：不投放，只更新物理系统
                game_over, victory = self.game.step(None)
            else:
                # 动作1-N：在指定位置投放
                position_idx = action - 1  # 减1因为动作0是特殊的"不投放"
                position = int(
                    (position_idx + 0.5) * self.width / self.discrete_positions
                )

                # 设置位置并调用投放
                self.game.drop_position = position
                self.game.drop_fruit()
                game_over, victory = self.game.step(None)
        else:
            # 处理None动作（只更新物理）
            game_over, victory = self.game.step(None)

        # 处理剩余的帧跳过
        for _ in range(1, self.frame_skip):
            if game_over:
                break  # 如果游戏结束，不再模拟额外帧

            # 只更新物理系统，不做新的输入
            next_game_over, next_victory = self.game.step(None)
            game_over = game_over or next_game_over
            victory = victory or next_victory

            # 计算中间帧奖励
            cumulative_reward += (self.game.score - self._last_score) / 100.0
            self._last_score = self.game.score

            # 如果启用渲染，也渲染中间帧
            if self.render_mode == "human":
                self._render_frame()

        # 获取最终观察
        observation = self._get_observation()

        # 计算最终奖励
        reward = cumulative_reward

        # 根据最终得分调整奖励
        score_diff = self.game.score - self._last_score
        reward += score_diff / 100.0
        self._last_score = self.game.score

        # 记录本次动作
        self._last_action = action

        # 检测胜利额外奖励
        if self.game.victory:
            reward += 10000.0  # 大奖励

        # 判断是否终止
        terminated = game_over
        truncated = False

        # 额外信息
        info = {"score": self.game.score, "victory": self.game.victory}

        # 渲染最终帧
        if self.render_mode == "human":
            self._render_frame()

        return observation, reward, terminated, truncated, info

    def render(self):
        """渲染当前帧"""
        if self.render_mode == "rgb_array":
            return self._get_observation()
        elif self.render_mode == "human":
            return self._render_frame()

    def _render_frame(self):
        """渲染帧到窗口"""
        if self.window is None and self.render_mode == "human":
            pygame.init()
            pygame.display.init()
            self.window = pygame.display.set_mode((self.width, self.height))
            pygame.display.set_caption("合成大西瓜")

        if self.clock is None and self.render_mode == "human":
            self.clock = pygame.time.Clock()

        # 获取游戏画面并显示
        canvas = self.game.render()

        if self.render_mode == "human":
            # 复制到窗口并刷新
            self.window.blit(canvas, canvas.get_rect())
            pygame.event.pump()
            pygame.display.update()

            # 控制帧率
            self.clock.tick(self.metadata["render_fps"])

    def close(self):
        """关闭环境"""
        if self.window is not None:
            pygame.display.quit()
            pygame.quit()
            self.window = None
            self.clock = None


class ProcessWrapper(mp.Process):
    """处理单个环境的工作进程"""

    def __init__(self, env_fn, pipe, seed=None):
        super(ProcessWrapper, self).__init__()
        self.env_fn = env_fn
        self.pipe = pipe
        self.seed = seed
        self.daemon = True  # 设置为守护进程，主进程退出时自动终止

    def preprocess_observation(self, observation):
        """在子进程中预处理观察"""
        # 转换为(C, H, W)格式并归一化
        return np.transpose(observation, (2, 0, 1)) / 255.0

    def run(self):
        """进程主循环"""
        env = self.env_fn()
        last_obs = None  # 保存最后一个有效观察

        if self.seed is not None:
            obs, _ = env.reset(seed=self.seed)
            last_obs = self.preprocess_observation(obs)  # 初始化最后观察

        while True:
            cmd, data = self.pipe.recv()

            if cmd == "step":
                obs, reward, terminated, truncated, info = env.step(data)

                if terminated:  # 当环境终止时
                    # 返回零填充的观察而不是None
                    zero_obs = np.zeros_like(last_obs) if last_obs is not None else None
                    self.pipe.send((zero_obs, reward, terminated, truncated, info))
                else:
                    # 处理正常情况
                    processed_obs = self.preprocess_observation(obs)
                    last_obs = processed_obs  # 更新最后一个有效观察
                    self.pipe.send((processed_obs, reward, terminated, truncated, info))

            elif cmd == "reset":
                obs, info = env.reset(seed=data if data else None)
                processed_obs = self.preprocess_observation(obs)
                last_obs = processed_obs  # 更新最后一个有效观察
                self.pipe.send((processed_obs, info))

            elif cmd == "close":
                env.close()
                self.pipe.close()
                break

            elif cmd == "render":
                self.pipe.send(env.render())

            else:
                raise NotImplementedError(f"Command {cmd} not implemented")


class MultiProcessVectorizedEnv:
    """使用多进程管理多个并行的西瓜环境"""

    def __init__(self, env_fn, num_envs=4, context="spawn"):
        """
        初始化多个并行环境

        参数:
            env_fn: 创建环境的函数
            num_envs: 并行环境数量
            context: 多进程上下文类型
        """
        ctx = mp.get_context(context)

        self.num_envs = num_envs
        self.env_fn = env_fn
        self.waiting = False  # 是否正在等待子进程

        # 为每个环境创建通信管道
        self.parent_pipes = []
        self.child_pipes = []

        for _ in range(num_envs):
            parent_conn, child_conn = ctx.Pipe()
            self.parent_pipes.append(parent_conn)
            self.child_pipes.append(child_conn)

        # 创建并启动子进程
        self.processes = []
        for i, child_pipe in enumerate(self.child_pipes):
            process = ProcessWrapper(env_fn, child_pipe, seed=i)
            self.processes.append(process)
            process.start()

        # 关闭父进程不需要的子进程连接端
        for child_pipe in self.child_pipes:
            child_pipe.close()

    def reset(self, seeds=None):
        """重置所有环境"""
        if seeds is None:
            seeds = [None] * self.num_envs

        for pipe, seed in zip(self.parent_pipes, seeds):
            pipe.send(("reset", seed))

        results = [pipe.recv() for pipe in self.parent_pipes]
        observations = np.array([res[0] for res in results])
        infos = [res[1] for res in results]

        return observations, infos

    # 添加到 MultiProcessVectorizedEnv 类中
    def reset_one(self, env_idx, seed=None):
        """重置单个环境"""
        self.parent_pipes[env_idx].send(("reset", seed))
        result = self.parent_pipes[env_idx].recv()
        return result  # 返回(observation, info)

    def step(self, actions):
        """在所有环境中执行动作"""
        # 发送动作到所有环境
        for pipe, action in zip(self.parent_pipes, actions):
            pipe.send(("step", action))

        # 接收所有环境的结果
        results = [pipe.recv() for pipe in self.parent_pipes]

        obs = np.array([res[0] for res in results])
        rews = np.array([res[1] for res in results])
        terms = np.array([res[2] for res in results])
        truncs = np.array([res[3] for res in results])
        infos = [res[4] for res in results]

        return obs, rews, terms, truncs, infos

    def close(self):
        """关闭所有环境和进程"""
        if self.processes:
            for pipe in self.parent_pipes:
                pipe.send(("close", None))

            for process in self.processes:
                process.join(timeout=5)  # 等待进程终止，最多5秒

            for pipe in self.parent_pipes:
                pipe.close()

            self.processes = []
            self.parent_pipes = []
